//------------------------------------------------------------------------------
// RegPH.cpp
//------------------------------------------------------------------------------
// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF 
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO 
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A 
// PARTICULAR PURPOSE.
//------------------------------------------------------------------------------
//
// DLL exports implementation for Protocol Handler
//
//------------------------------------------------------------------------------

#include "pch.h"
#include <msi.h>
#include <msidefs.h>
#pragma warning(push)
#pragma warning(disable:4995)
#include <atlsecurity.h>
#pragma warning(pop)
#include <Utils.h>
#include <RegNamespace.h>

class CRegPHModule : public CAtlDllModuleT<CRegPHModule>
{
public :
	DECLARE_REGISTRY_APPID_RESOURCEID(IDR_REGPH, "{70B74D2D-DE5A-42FD-876B-197803BB4F31}")

    HRESULT RegisterServer(BOOL fTypeLib);
    HRESULT UnregisterServer(BOOL fTypeLib);
    HRESULT CreateDeleteRegSearchRootAndDefaultScope(BOOL fCreate, PCWSTR pszUserSID);

private:
    HRESULT _CreateSearchRootAndDefaultScope(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL);
    HRESULT _DeleteSearchRoot(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL);
    HRESULT _FindSearchRoot(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL, __out BOOL *pfFoundRoot);
    HRESULT _GetUserSIDString(__deref_out PWSTR *ppszSID);
};

CRegPHModule _AtlModule;

HRESULT CRegPHModule::_GetUserSIDString(__deref_out PWSTR *ppszSID)
{
    CAccessToken tkn;
    HRESULT hr = E_FAIL;
    if (tkn.GetEffectiveToken(TOKEN_QUERY))
    {
        CSid sid;
        try
        {
            if (tkn.GetUser(&sid))
            {
                DWORD cchSID = lstrlen(sid.Sid()) + 1;
                CComMemPtr<WCHAR> spszSID;
                hr = spszSID.Allocate(cchSID);
                if (SUCCEEDED(hr))
                {
                    hr = StringCchCopy(spszSID, cchSID, sid.Sid());
                    if (SUCCEEDED(hr))
                    {
                        *ppszSID = spszSID.Detach();
                    }
                }
            }
        }
        catch (CAtlException &e)
        {
            hr = e.m_hr;
        }
    }
    return hr;
}

HRESULT CRegPHModule::RegisterServer(BOOL fTypeLib)
{
    HRESULT hr = CAtlDllModuleT::RegisterServer(fTypeLib);
#ifdef _MERGE_PROXYSTUB
    if (SUCCEEDED(hr))
    {
        hr = PrxDllRegisterServer();
    }
#endif
    if (SUCCEEDED(hr))
    {
        CComMemPtr<WCHAR> spszSID;
        hr = _GetUserSIDString(&spszSID);
        if (SUCCEEDED(hr))
        {
            hr = CreateDeleteRegSearchRootAndDefaultScope(TRUE, spszSID);
        }
    }
    return hr;
}

HRESULT CRegPHModule::UnregisterServer(BOOL fTypeLib)
{
    CComMemPtr<WCHAR> spszSID;
    HRESULT hr = _GetUserSIDString(&spszSID);
    if (SUCCEEDED(hr))
    {
        hr = CreateDeleteRegSearchRootAndDefaultScope(FALSE, spszSID);
        if (SUCCEEDED(hr))
        {
            hr = CAtlDllModuleT::UnregisterServer(fTypeLib);
#ifdef _MERGE_PROXYSTUB
            if (SUCCEEDED(hr))
            {
                hr = PrxDllRegisterServer();
                if (SUCCEEDED(hr))
                {
                    hr = PrxDllUnregisterServer();
                }
            }
#endif
        }
    }
    return hr;
}

HRESULT CRegPHModule::_FindSearchRoot(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL, __out BOOL *pfFoundRoot)
{
    *pfFoundRoot = FALSE;
    CComPtr<IEnumSearchRoots> spesr;
    HRESULT hr = pscsm->EnumerateRoots(&spesr);
    if (hr == S_OK)
    {
        // Check if has this search root already
        CComPtr<ISearchRoot> spsr;
        while ((hr = spesr->Next(1, &spsr, NULL)) == S_OK && !*pfFoundRoot)
        {
            CComMemPtr<WCHAR> spszRootURL;
            hr = spsr->get_RootURL(&spszRootURL);
            if (SUCCEEDED(hr))
            {
                // Check if this is our reg://{SID}/ URL
                if (CompareString(LOCALE_INVARIANT, NORM_IGNORECASE, pszURL, -1, spszRootURL, -1) == CSTR_EQUAL)
                {
                    // Found a dupe, hr is S_OK
                    *pfFoundRoot = TRUE;
                }
            }

            // Looping on ISearchResult, release it here
            spsr.Release();
        }
    }
    return SUCCEEDED(hr) ? S_OK : hr;
}

HRESULT CRegPHModule::_DeleteSearchRoot(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL)
{
    HRESULT hr = pscsm->RemoveRoot(pszURL);
    if (SUCCEEDED(hr))
    { 
        hr = pscsm->SaveAll();
    }
    return hr;
}

HRESULT CRegPHModule::_CreateSearchRootAndDefaultScope(ISearchCrawlScopeManager *pscsm, PCWSTR pszURL)
{
    CComPtr<ISearchRoot> spsr;
    HRESULT hr = spsr.CoCreateInstance(__uuidof(CSearchRoot), NULL, CLSCTX_SERVER);
    if (SUCCEEDED(hr))
    {
        // Need this so that the WDS API's work even though the user is not an Admin
        hr = CoSetProxyBlanket(spsr, RPC_C_AUTHN_DEFAULT, RPC_C_AUTHZ_DEFAULT,
            NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);
        if (SUCCEEDED(hr))
        {
            hr = spsr->put_RootURL(pszURL);
            if (SUCCEEDED(hr))
            {
                hr = spsr->put_ProvidesNotifications(FALSE);
                if (SUCCEEDED(hr))
                {
                    hr = spsr->put_IsHierarchical(TRUE);
                    if (SUCCEEDED(hr))
                    {
                        hr = spsr->put_UseNotificationsOnly(FALSE);
                        if (SUCCEEDED(hr))
                        {
                            hr = spsr->put_FollowDirectories(TRUE);
                            if (SUCCEEDED(hr))
                            {
                                hr = pscsm->AddRoot(spsr);
                                if (SUCCEEDED(hr))
                                {
                                    hr = pscsm->AddHierarchicalScope(pszURL, TRUE, TRUE, FALSE);
                                    if (SUCCEEDED(hr))
                                    {
                                        hr = pscsm->SaveAll();
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return hr;
}

HRESULT CRegPHModule::CreateDeleteRegSearchRootAndDefaultScope(BOOL fCreate, PCWSTR pszUserSID)
{
    CComPtr<ISearchManager> spsm;
    HRESULT hr = spsm.CoCreateInstance(__uuidof(CSearchManager), NULL, CLSCTX_SERVER);
    if (SUCCEEDED(hr))
    {
        CComPtr<ISearchCatalogManager> spscm;
        hr = spsm->GetCatalog(L"SystemIndex", &spscm);
        if (SUCCEEDED(hr))
        { 
            CComPtr<ISearchCrawlScopeManager> spscsm;
            hr = spscm->GetCrawlScopeManager(&spscsm);
            if (SUCCEEDED(hr))
            {
                // Format the default URL for our Res protocol
                CBufString strURL;
                hr = StringCchPrintf(strURL, strURL.Size(), L"reg://{%s}/", pszUserSID);
                if (SUCCEEDED(hr))
                {
                    BOOL fFoundRoot;
                    hr = _FindSearchRoot(spscsm, strURL, &fFoundRoot);
                    if (SUCCEEDED(hr))
                    {
                        if (fCreate)
                        {
                            // Do an explicit remove to start with a clean MSSearch registration
                            if (fFoundRoot)
                            {
                                _DeleteSearchRoot(spscsm, strURL);
                            }

                            // Need to add the root and default scope
                            hr = _CreateSearchRootAndDefaultScope(spscsm, strURL);
                        }
                        else if (fFoundRoot)
                        {
                            // Need to remove the search root
                            hr = _DeleteSearchRoot(spscsm, strURL);
                        }
                    }
                }
            }
        }
    }
    return hr;
}

// DLL Entry Point
extern "C" BOOL WINAPI DllMain(HINSTANCE hInstance, DWORD dwReason, LPVOID lpReserved)
{
#ifdef _MERGE_PROXYSTUB
    if (!PrxDllMain(hInstance, dwReason, lpReserved))
        return FALSE;
#endif
	hInstance;
    return _AtlModule.DllMain(dwReason, lpReserved); 
}


// Used to determine whether the DLL can be unloaded by OLE
STDAPI DllCanUnloadNow(void)
{
#ifdef _MERGE_PROXYSTUB
    HRESULT hr = PrxDllCanUnloadNow();
    if (hr != S_OK)
        return hr;
#endif
    return _AtlModule.DllCanUnloadNow();
}


// Returns a class factory to create an object of the requested type
STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
{
#ifdef _MERGE_PROXYSTUB
    if (PrxDllGetClassObject(rclsid, riid, ppv) == S_OK)
        return S_OK;
#endif
    return _AtlModule.DllGetClassObject(rclsid, riid, ppv);
}

// DllRegisterServer - Adds entries to the system registry
STDAPI DllRegisterServer(void)
{
    return _AtlModule.DllRegisterServer(FALSE);
}


// DllUnregisterServer - Removes entries from the system registry
STDAPI DllUnregisterServer(void)
{
    return _AtlModule.DllUnregisterServer(FALSE);
}

extern "C" UINT WINAPI ManageRegistryStartAddress(MSIHANDLE hmsi)
{
    CComMemPtr<WCHAR> spszCustomActionData;
    HRESULT hr = GetMsiProperty(hmsi, IPROPNAME_CUSTOMACTIONDATA, &spszCustomActionData);
    if (SUCCEEDED(hr))
    {
        CComMemPtr<WCHAR> spszUserSID;
        hr = GetMsiProperty(hmsi, IPROPNAME_USERSID, &spszUserSID);
        if (SUCCEEDED(hr))
        {
            if (StrCmpIC(spszCustomActionData, L"Install") == 0)
            {
                hr = _AtlModule.CreateDeleteRegSearchRootAndDefaultScope(TRUE, spszUserSID);
            }
            else
            {
                hr = _AtlModule.CreateDeleteRegSearchRootAndDefaultScope(FALSE, spszUserSID);
            }
        }
    }
    return SCODE_CODE(GetScode(hr));
}
